# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Aggregator for node vector
"""

from typing import Union, List, Optional

import torch
from torch import nn, Tensor

from ..utils import get_arguments, GLOBAL_DEVICE
from ..layer import MLP, Dense

__all__ = [
    "InteractionAggregator",
    "InteractionSummation",
    "InteractionMean",
    "LinearTransformation",
    "MultipleChannelRepresentation",
    "get_interaction_aggregator",
]

_INTERACTION_AGGREGATOR_BY_KEY = {}


def _interaction_aggregator_register(*aliases):
    """Return the alias register."""
    def alias_reg(cls):
        name = cls.__name__.lower()
        if name not in _INTERACTION_AGGREGATOR_BY_KEY:
            _INTERACTION_AGGREGATOR_BY_KEY[name] = cls

        for alias in aliases:
            if alias not in _INTERACTION_AGGREGATOR_BY_KEY:
                _INTERACTION_AGGREGATOR_BY_KEY[alias] = cls

        return cls

    return alias_reg


class InteractionAggregator(nn.Module):
    r"""Network to aggregate the representation of each interaction layer.

    Args:
        dim (int):          Feature dimension.
        num_agg (int):      Number of interaction layer to be aggregate. Default: None
        n_hidden (int):     Number of hidden layers. Default: 0
        activation (nn.Module):  Activation function. Default: None

    Symbols:
        B:  Number of simulation walker.
        A:  Number of atoms in system.
        F:  Feature dimension.
    """

    def __init__(self, **kwargs):
        super().__init__()
        self._kwargs = kwargs
        self.device = GLOBAL_DEVICE()

    def forward(self,
                ylist: List[torch.Tensor],
                atom_mask: Optional[torch.Tensor] = None):
        """Aggregate the representations of each interaction layer.

        Args:
            ylist (list):       List of representation of interactions layers.
            atom_mask (Tensor): Tensor of shape (B, A). Data type is bool.
                                Mask for atoms.Default: None

        Returns:
            output (Tensor):    Tensor of shape (B, X). Data type is float.
        """
        raise NotImplementedError


@_interaction_aggregator_register('sum')
class InteractionSummation(InteractionAggregator):
    r"""A interaction aggregator to summation all representations of interaction layers

    Args:
        dim (int):          Feature dimension.
        num_agg (int):      Number of interaction layer to be aggregate. Default: None
        n_hidden (int):     Number of hidden layers. Default: 0
        activation (nn.Module):  Activation function. Default: None

    Symbols:
        B:  Number of simulation walker.
        A:  Number of atoms in system.
        F:  Feature dimension.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._kwargs = get_arguments(locals(), kwargs)
        self.reg_key = 'sum'

    def __str__(self):
        return "sum"

    def forward(self,
                ylist: List[torch.Tensor],
                atom_mask: Optional[torch.Tensor] = None):
        xt = torch.stack(ylist, dim=-1)
        y = torch.sum(xt, dim=-1)
        if atom_mask is not None:
            y = y * torch.unsqueeze(atom_mask, -1)
        return y


@_interaction_aggregator_register('mean')
class InteractionMean(InteractionAggregator):
    r"""A interaction aggregator to average all representations of interaction layers

    Args:
        dim (int):          Feature dimension.
        num_agg (int):      Number of interaction layer to be aggregate. Default: None
        n_hidden (int):     Number of hidden layers. Default: 0
        activation (nn.Module):  Activation function. Default: None

    Symbols:
        B:  Number of simulation walker.
        A:  Number of atoms in system.
        F:  Feature dimension.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._kwargs = get_arguments(locals(), kwargs)
        self.reg_key = 'mean'

    def __str__(self):
        return "mean"

    def forward(self,
                ylist: List[torch.Tensor],
                atom_mask: Optional[torch.Tensor] = None):
        xt = torch.stack(ylist, dim=-1)
        y = torch.mean(xt, dim=-1)
        if atom_mask is not None:
            y = y * torch.unsqueeze(atom_mask, -1)
        return y


@_interaction_aggregator_register('linear')
class LinearTransformation(InteractionAggregator):
    r"""A interaction aggregator to aggregate all representations of interaction layers
        by using linear transformation

    Args:
        dim (int):          Feature dimension.
        num_agg (int):      Number of interaction layer to be aggregate. Default: None
        n_hidden (int):     Number of hidden layers. Default: 0
        activation (nn.Module):  Activation function. Default: None

    Symbols:
        B:  Number of simulation walker.
        A:  Number of atoms in system.
        F:  Feature dimension.
    """

    def __init__(self, dim: int, **kwargs):
        super().__init__(**kwargs)
        self._kwargs = get_arguments(locals(), kwargs)
        self.reg_key = 'linear'
        self.dim = dim

        # Initialize parameters
        self.scale = nn.Parameter(torch.randn(dim, device=self.device))
        self.shift = nn.Parameter(torch.randn(dim, device=self.device))

    def __str__(self):
        return "linear"

    def forward(self,
                ylist: List[torch.Tensor],
                atom_mask: Optional[torch.Tensor] = None):
        yt = torch.stack(ylist, dim=-1)
        ysum = torch.sum(yt, dim=-1)
        y = self.scale * ysum + self.shift
        if atom_mask is not None:
            y = y * torch.unsqueeze(atom_mask, -1)
        return y


@_interaction_aggregator_register('mcr')
class MultipleChannelRepresentation(InteractionAggregator):
    r"""A Multiple-Channel Representation (MCR) interaction aggregator to
        aggregate all representations of interaction layers

    Args:
        dim (int):          Feature dimension.
        num_agg (int):      Number of interaction layer to be aggregate.
        n_hidden (int):     Number of hidden layers. Default: 0
        activation (nn.Module):  Activation function. Default: None

    Symbols:
        B:  Number of simulation walker.
        A:  Number of atoms in system.
        F:  Feature dimension.
    """

    def __init__(self,
                 dim: int,
                 num_agg: int,
                 n_hidden: int = 0,
                 activation: Optional[nn.Module] = None,
                 **kwargs,
                 ):

        super().__init__(**kwargs)
        self._kwargs = get_arguments(locals(), kwargs)

        self.reg_key = 'mcr'
        self.dim = dim
        self.num_agg = num_agg
        self.n_hidden = n_hidden
        self.activation = activation

        sub_dim = self.dim // self.num_agg
        last_dim = self.dim - (sub_dim * (self.num_agg - 1))
        sub_dims = [sub_dim for _ in range(self.num_agg - 1)]
        sub_dims.append(last_dim)

        if self.n_hidden > 0:
            hidden_layers = [dim] * self.n_hidden
            self.mcr = nn.ModuleList([
                MLP(
                    self.dim, 
                    sub_dims[i], 
                    hidden_layers,
                    activation=self.activation,
                )
                for i in range(self.num_agg)
            ])
        else:
            self.mcr = nn.ModuleList([
                Dense(self.dim, sub_dims[i])
                for i in range(self.num_agg)
            ])

    def __str__(self):
        return "MCR"

    def forward(self,
                ylist: List[torch.Tensor],
                atom_mask: Optional[torch.Tensor] = None):
        readouts = []
        for i in range(self.num_agg):
            readouts.append(self.mcr[i](ylist[i]))
        y = torch.cat(readouts, dim=-1)
        if atom_mask is not None:
            y = y * torch.unsqueeze(atom_mask, -1)
        return y


_INTERACTION_AGGREGATOR_BY_NAME = {
    agg.__name__: agg for agg in _INTERACTION_AGGREGATOR_BY_KEY.values()
}


def get_interaction_aggregator(cls_name: Union[InteractionAggregator, str, dict],
                               **kwargs,
                               ) -> InteractionAggregator:
    """get aggregator by name"""
    if cls_name is None or isinstance(cls_name, InteractionAggregator):
        return cls_name
    if isinstance(cls_name, dict):
        return get_interaction_aggregator(**cls_name)
    if isinstance(cls_name, str):
        if cls_name.lower() == 'none':
            return None
        if cls_name.lower() in _INTERACTION_AGGREGATOR_BY_KEY.keys():
            return _INTERACTION_AGGREGATOR_BY_KEY[cls_name.lower()](**kwargs)
        if cls_name in _INTERACTION_AGGREGATOR_BY_NAME.keys():
            return _INTERACTION_AGGREGATOR_BY_NAME[cls_name](**kwargs)
        raise ValueError(
            "The Interaction Aggregator corresponding to '{}' was not found.".format(cls_name))
    raise TypeError(
        "Unsupported Interaction Aggregator type '{}'.".format(type(cls_name)))
